import tensorflow as tf


def conv_block(inp, filters):
    x = tf.keras.layers.Conv2D(filters, (3, 3), padding='same', activation='relu')(inp)
    x = tf.keras.layers.Conv2D(filters, (3, 3), padding='same')(x)
    x = tf.keras.layers.BatchNormalization(axis=3)(x)
    x = tf.keras.layers.Activation('relu')(x)
    return x


def encoder_block(inp, filters, dropout):
    x = conv_block(inp, filters)
    p = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
    p = tf.keras.layers.Dropout(dropout)(p)
    return x, p


def attention_block(l_layer, h_layer):  # Attention Block
    phi = tf.keras.layers.Conv2D(h_layer.shape[-1], (1, 1), padding='same')(l_layer)
    theta = tf.keras.layers.Conv2D(h_layer.shape[-1], (1, 1), strides=(2, 2), padding='same')(h_layer)
    x = tf.keras.layers.add([phi, theta])
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.Conv2D(1, (1, 1), padding='same', activation='sigmoid')(x)

    x = tf.keras.layers.UpSampling2D(size=(2, 2))(x)
    x = tf.keras.layers.multiply([h_layer, x])
    x = tf.keras.layers.BatchNormalization(axis=3)(x)
    return x


def decoder_block(inp, filters, concat_layer, dropout):
    x = tf.keras.layers.Conv2DTranspose(filters, (2, 2), strides=(2, 2), padding='same')(inp)
    concat_layer = attention_block(inp, concat_layer)
    x = tf.keras.layers.concatenate([x, concat_layer])
    x = tf.keras.layers.Dropout(dropout)(x)
    x = conv_block(x, filters)
    return x


def create_model(input_shape):
    input_img = tf.keras.Input(input_shape, name='img')
    d1, p1 = encoder_block(input_img, 64, 0.1)
    d2, p2 = encoder_block(p1, 128, 0.1)
    d3, p3 = encoder_block(p2, 256, 0.1)
    d4, p4 = encoder_block(p3, 512, 0.1)
    b1 = conv_block(p4, 1024)
    e2 = decoder_block(b1, 512, d4, 0.1)
    e3 = decoder_block(e2, 256, d3, 0.1)
    e4 = decoder_block(e3, 128, d2, 0.1)
    e5 = decoder_block(e4, 64, d1, 0.1)
    outputs = tf.keras.layers.Conv2D(1, (1, 1), activation="sigmoid")(e5)
    model = tf.keras.Model(inputs=[input_img], outputs=[outputs], name='AttentionUnet')
    return model
